from types import SimpleNamespace

from wenet.zipformer_new.decoder_attention import AttentionDecoderModel

def init_decoder4zipformer(config_zipformer):
    params = SimpleNamespace(**config_zipformer)
    decoder = AttentionDecoderModel(
        vocab_size=params.vocab_size,
        decoder_dim=params.attention_decoder_dim,
        num_decoder_layers=params.attention_decoder_num_layers,
        attention_dim=params.attention_decoder_attention_dim,
        num_heads=int(params.attention_decoder_num_heads),
        feedforward_dim=params.attention_decoder_feedforward_dim,
        memory_dim=max(params.encoder_dim),
        sos_id=params.sos_id,
        eos_id=params.eos_id,
        ignore_id=params.ignore_id,
        label_smoothing=params.label_smoothing,
    )
    return decoder

